#!/usr/bin/env python3
"""
clean4llm.py - Git仓库清洗工具 for LLM知识库
功能：复制代码文件时排除二进制/媒体/构建产物等非必要文件
"""

import os
import sys
import fnmatch
import shutil
import subprocess
import argparse
from pathlib import Path
from tqdm import tqdm
import time
from collections import defaultdict
from wcwidth import wcswidth

# 排除规则配置（可扩展）
FILE_EXCLUDES = {
    'extensions': {
    # 强制忽略的二进制文件类型
    '.exe', '.bin', '.so', '.pyd',
    '.doc', '.xls', '.ppt',
    '.a', '.lib', '.o', '.obj',
    '.class', '.prefs',
    # 强制忽略 java 打包文件
    '.jar', '.war', '.ear',
    # 强制忽略压缩包
    '.zip', '.rar', '.7z', '.bz2' '.iso', '.tar', '.gz', '.tgz',
    '.xz', '.z', '.arj', '.cab', '.lzh', '.ex_', '.dl_',
    '.log', 
    # 强制忽略图像格式文件
    '.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.tif', '.webp', '.psd', '.raw',
    '.heic', '.pcx', '.tga', '.jp2', '.j2k', '.iff', '.fpx', 
    # 强制忽略矢量图像格式
    '.svg', '.ai', '.eps', 
    '.cdr', '.dwg', '.dxf', '.ico', '.xbm', '.xpm',
    # 强制忽略医学专用图像文件
    '.dcm', '.nii', '.img', '.hdr', '.nrrd', '.mnc', '.par', '.rec', '.svs',
    # 强制忽略视频格式文件
    '.mp4', '.avi', '.mkv', '.mov', '.flv', '.wmv', '.rmvb', '.3gp', 
    '.webm', '.ts', '.m2ts', '.qlv', '.vob', '.asf',
    # 强制忽略音频格式文件
    '.wav', '.mp3', '.flac', '.aac', '*aif' '.aiff', '.ape', '.ogg', '.wma',
    '.mid', '.midi', '.cda', '.ra', '.rm', '.au', '.vqf', '.dsf', '.dff',
    # 强制忽略数据库文件
    '.rdb', '.aof', '.db', '.mdf', '.ndf', '.ldf', '.frm', '.ibd', '.myd', '.myi', '.dbf',
    '.ctl', '.dmp', '.dat', '.backup', '.sqlite', '.mdb', '.accdb', '.bson', '.ibdata',
    '.fdb', '.hdb', '.backup',
    # 强制忽略二进制模型文件
    '.gguf', '.pt', '.pth', '.ckpt', '.onnx', '*.safetensors'
    },
    'patterns': ['.*', '~*']  # 隐藏文件和临时文件
}

DIR_EXCLUDES = {
    'names': {'bin', 'build', 'target', 'lib', 'log', 'node_modules'},
    'patterns': ['.git', '.*/']
}

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='清理Git仓库用于LLM知识库导入',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument('source', help='源目录路径')
    parser.add_argument('--output', help='目标目录路径')
    parser.add_argument("-v",'--verbose', action='store_true', help='显示详细输出')
    return parser.parse_args()

def should_exclude_dir(dir_path, source_root):
    """判断是否排除目录"""
    dir_name = os.path.basename(dir_path)
    
    # 检查目录名黑名单
    if dir_name in DIR_EXCLUDES['names']:
        return True
    
    # 检查路径模式
    rel_path = os.path.relpath(dir_path, source_root)
    for pattern in DIR_EXCLUDES['patterns']:
        if fnmatch.fnmatch(rel_path, pattern):
            return True
    return False

def should_exclude_file(file_path, source_root):
    """判断是否排除文件"""
    # 检查文件扩展名
    ext = os.path.splitext(file_path)[1].lower()
    if ext in FILE_EXCLUDES['extensions']:
        return True
    
    # 检查文件模式
    rel_path = os.path.relpath(file_path, source_root)
    for pattern in FILE_EXCLUDES['patterns']:
        if fnmatch.fnmatch(rel_path, pattern):
            return True
    
    # 检查.gitignore规则
    if check_git_ignore(source_root, rel_path):
        return True
    
    return False

def check_git_ignore(repo_root, rel_path):
    """检查文件是否被.gitignore排除"""
    try:
        result = subprocess.run(
            ['git', '-C', repo_root, 'check-ignore', '-q', rel_path],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
        return result.returncode == 0
    except Exception as e:
        print(f"Git检查错误: {e}", file=sys.stderr)
        return False

def copy_clean_files(src, dst, verbose=False):
    """执行清洗复制操作，返回(总文件数, 成功数)"""
    from collections import defaultdict
    from tqdm import tqdm
    print("\033[34m统计文件数量。。。\033[0m")
    # 预生成有效文件列表
    file_list = []
    for root, dirs, files in os.walk(src):
        dirs[:] = [d for d in dirs if not should_exclude_dir(os.path.join(root, d), src)]
        for file in files:
            src_path = os.path.join(root, file)
            if not should_exclude_file(src_path, src):
                file_list.append((root, file))

    total = len(file_list)
    if total == 0:
        print("没有发现需要复制的文件")
        return 0

    if verbose:
        print("\n\033[33m要复制的文件列表：\033[0m")
        for root, file in file_list:
            print(f"- {os.path.join(root, file)}")

    print(f"开始复制 {src} 的 {total} 个文件")
    file_stats = defaultdict(int)

    with tqdm(total=total, desc="处理进度", unit="文件", 
             bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]") as pbar:
        for root, file in file_list:
            src_path = os.path.join(root, file)
            rel_path = os.path.relpath(src_path, src)
            dst_path = os.path.join(dst, rel_path)
            
            try:
                # 创建目录结构
                os.makedirs(os.path.dirname(dst_path), exist_ok=True)
                
                # 执行复制
                shutil.copy2(src_path, dst_path)

                # 统计文件类型
                key = generate_stat_key(file)
                file_stats[key] += 1

            except Exception as e:
                print(f"复制失败: {rel_path} - {e}", file=sys.stderr)
            finally:
                pbar.update(1)
                pbar.set_postfix(file=truncate_path(rel_path, 30))

    # echo -e "\033[32m复制完成：源目录已复制到\033[0m $TARGET_DIR"
    print(f"\033[32m复制完成 源目录已复制到\033[0m {dst}")
    # 输出统计信息
    print("\n===== 统计信息 =====")
    if file_stats.keys():
        # 计算动态列宽
        max_key_display_width = max(get_display_width(ext) for ext in file_stats.keys()) if file_stats else 0 
        col1_width = max(max_key_display_width, 12)  # 基于显示宽度计算 [^1]
        col2_width = 8

        # 生成表格
        table_header = (
            f"┌{'─'*(col1_width+2)}┬{'─'*(col2_width+2)}┐\n"
            f"│ {format_cell('类型', col1_width)} │ {format_cell('数量', col2_width)} │\n"  # 使用对齐函数
            f"├{'─'*(col1_width+2)}┼{'─'*(col2_width+2)}┤"
        )
        print(table_header)

        for ext, count in sorted(file_stats.items(), key=lambda x: x[1], reverse=True):
            print(
                f"│ {format_cell(ext, col1_width)} │ {format_cell(count, col2_width)} │"  # 统一使用对齐函数
            )

        print(f"└{'─'*(col1_width+2)}┴{'─'*(col2_width+2)}┘")
    else:
        print("无")

def generate_stat_key(file_path: str) -> str:
    """生成标准化的文件类型统计键"""
    filename = Path(file_path).name
    
    # 处理无扩展名的独立文件（如Dockerfile、LICENSE）
    if '.' not in filename or filename.startswith('.'):
        return filename
    
    # 处理带扩展名的文件（如main.py -> *.py）
    return f"*.{filename.split('.')[-1].lower()}"

def truncate_path(path, max_len=40):
    """路径折叠显示实现"""
    if len(path) <= max_len:
        return path
    return path[:max_len//3] + "..." + path[-2*max_len//3:]

def get_display_width(text):
    return wcswidth(str(text))  # 准确计算中文显示宽度 [^1]

def format_cell(content, width):
    """带中文字符对齐的格式化"""
    fill_size = width - (get_display_width(content) - len(str(content))) 
    return f"{str(content):<{fill_size}}"  # 修正填充逻辑 [^1]

def main():
    args = parse_args()
    source = os.path.abspath(args.source)
    
    if not os.path.isdir(source):
        print(f"错误：源目录不存在: {source}", file=sys.stderr)
        sys.exit(1)
    
    target = args.output or f"{os.path.basename(source)}_src"
    target = os.path.abspath(target)
    
    # 清空目标目录
    if os.path.exists(target):
        shutil.rmtree(target)
    os.makedirs(target, exist_ok=True)
    
    start_time = time.perf_counter()
    copy_clean_files(source, target, args.verbose)
    print(f"耗时: {time.perf_counter() - start_time:.2f}秒")
if __name__ == '__main__':
    main()
